# coding=utf-8

import numpy as np
import pylab as pl
from scipy.optimize import leastsq

# 线性回归

n = 2  # (多项式次数-1)或者是多项式参数个数, y=kx+b，参数个数为2，分别为k、b


# 多项式函数(拟合函数、目标函数，要求出的函数)
def fit_func(p, x):
    f = np.poly1d(p)
    return f(x)


# 残差函数（误差函数）
def residuals_func(p, y, x):
    ret = fit_func(p, x) - y
    return ret


# 数据
x = [1, 2, 3, 5, 9]
y = [0.9, 1.8, 3.2, 5.1, 8.5]

# 随机初始化多项式参数
p_init = np.random.randn(n)
param = leastsq(residuals_func, p_init, args=(y, x))

# 输出拟合参数
print('Fitting Param: ', param[0])
# Fitting Parameters:  [ -2.83923076e+02   1.14996949e+03  -2.06511905e+03   2.08842363e+03 -1.22138657e+03   4.03643606e+02  -8.37576650e+01   1.21696134e+01 8.60284229e-03]

# 绘制拟合函数
x_points = np.linspace(0, 10, 10000)
pl.plot(x_points, fit_func(param[0], x_points), label='fitted curve')

# 绘制
pl.plot(x, y, 'bo', label='data', color='red')

pl.legend()
pl.show()
